import datetime

import arrow
from flask import current_app

from app.db.models.vulndb import KbVuln
from app.db.models.balances import Balances
from app.db.models.tasks import Task
from app.db.models.asset_tasks import AssetTasks, get_monitor_urls
from app.db.models.jobs import Job
from app.db.models.alertreceivers import AlertReceivers
from app.db.models.alert_groups import AlertGroup
from app.libs.score import Score
from app.libs.report import keyword_category_mapping, privacy_category_mapping
from app.libs.display import WARNING_CATEGORY_KEY_LABEL_MAP


class ReportPushData(object):

    web_sec_type = ["content", "securityEvent", "ssl", "vul"]
    available_type = ["http", "ping"]
    level_num_en = {
        1: "low",
        2: "low",
        3: "medium",
        4: "high",
        5: "critical",
    }

    def __init__(self, user, start_time, end_time, kb_data: dict = None):
        self.user = user
        self.start_time = start_time
        self.end_time = end_time
        self.time_query = {"$gte": self.start_time, "$lte": self.end_time}
        self.balance = Balances.objects.filter(uid=self.user.id).first()
        self.kb_data = self.get_kb_data(kb_data)

    @staticmethod
    def utc_to_local(t):
        return arrow.get(t).to("local").datetime

    def get_receiver(self):
        receiver_data = {}
        if receiver := AlertReceivers.objects.filter(uid=self.user.id, isDefault=True).first():
            if receiver.phoneVerified and receiver.phone:
                receiver_data.update({"phone": [receiver.phone, ]})
            if receiver.emailVerified and receiver.email:
                receiver_data.update({"email": [receiver.email, ]})
        if alert_group := AlertGroup.objects.filter(uid=self.user.id, isDefault=True).first():
            if bool(alert_group.alertSettings.notification.get('wechat')) and self.user.wechat and (wechat := self.user.wechat.get("wx_openid")):
                receiver_data.update({"wechat": [wechat, ]})
        return receiver_data

    @property
    def is_free_user(self):
        """是否是付费用户"""
        if not self.balance:
            return True
        if not self.balance.planExpireTime:
            return True
        if arrow.get(self.balance.planExpireTime).datetime < self.start_time:
            return True
        return False

    def get_kb_data(self, kb_data):
        """漏洞更新数据"""
        if kb_data:
            return kb_data
        query = {'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}, 'is_deleted': False}
        all_vul = KbVuln.objects.find(
            query).only("source", "id", "updated_at", "severity", "vuln_type", "name")
        total = all_vul.count()
        this_time_data = all_vul.find({"updated_at": self.time_query})
        result = []
        exists_ids = []
        if this_time_data:
            for d in this_time_data.order_by("-updated_at", "-severity").limit(5):
                exists_ids.append(d.id)
                result.append(
                    {
                        "name": d.name.zh_cn,
                        "type": d.vuln_type[0].zh_cn if d.vuln_type else '',
                        "level": Score.severityMap[d.severity]["zh_CN"],
                        "severity": d.severity
                    }
                )
        if len(result) < 5:
            for d in all_vul.find({"_id": {"$nin": exists_ids}}).order_by("-updated_at", "-severity").limit(5-len(result)):
                result.append(
                    {
                        "name": d.name.zh_cn,
                        "type": d.vuln_type[0].zh_cn if d.vuln_type else '',
                        "level": Score.severityMap[d.severity]["zh_CN"],
                        "severity": d.severity
                    }
                )
        result = list(sorted(result, key=lambda x: x["severity"], reverse=True))
        return {"count": total, "detail": result}

    def pkg_data(self):
        """套餐到期时间， 套餐可接入资产数量"""
        used_domain_count = Job.objects.find({"uid": self.user.id, "assetType": {"$ne": "host"}}).count()
        free_count = 0
        expire_time = ""
        if self.balance:
            domain_total = self.balance.balance.domain.total
            free_count = domain_total - used_domain_count
        if self.balance and self.balance.planExpireTime:
            expire_time = self.utc_to_local(self.balance.planExpireTime).strftime("%Y-%m-%d")
        free_count = free_count if free_count > 0 else 0
        return used_domain_count, free_count, expire_time

    @property
    def duration_time(self):
        """服务持续时间"""
        return (arrow.utcnow().datetime - arrow.get(self.user.createAt).datetime).days or 1

    def asset_data(self):
        """只统计web资产， 不统计主机资产"""
        monitor_asset_query = {
            "uid": self.user.id,
            "taskGroupType": {"$in": ['web_security', 'http', 'ping']},
            "$or": [
                {"enableMonitor": True, "triggerType": "schedule"},
                {"triggerType": "manual", "status": {"$in": ["waiting", "active"]}}
            ]
        }
        monitor_asset_ids = AssetTasks.objects.find(monitor_asset_query).distinct("jobId")
        monitor_asset_count = len(monitor_asset_ids)
        not_monitor_asset_count_q = {"_id": {"$nin": monitor_asset_ids}, "assetType": {"$ne": "host"}, "uid": self.user.id}
        not_monitor_asset_count = Job.objects.find(not_monitor_asset_count_q).count()
        not_verified_asset_count_q = {"verification.isVerified": False, "assetType": {"$ne": "host"}, "uid": self.user.id}
        not_verified_asset_count = Job.objects.find(not_verified_asset_count_q).count()
        return monitor_asset_count, not_monitor_asset_count, not_verified_asset_count

    def asset_task_data(self):
        """资产执行任务次数，结果异常任务资产数"""
        task_base_query = {"uid": self.user.id, "endTime": self.time_query, "status": "completed"}
        sec_q = {"taskType": {"$in": self.web_sec_type}, "securityStatus": "warning", **task_base_query}
        sec_risk_asset_count = len(Task.objects.find(sec_q).distinct("jobId"))
        ava_q = {"taskType": {"$in": self.available_type}, "securityStatus": "warning", **task_base_query}
        available_risk_asset_ids = Task.objects.find(ava_q).distinct("jobId")
        available_risk_asset_count = len(available_risk_asset_ids)
        safe_ava_q = {"taskType": {"$in": self.available_type}, "securityStatus": "safe",
                      "jobId": {"$nin": available_risk_asset_ids}, **task_base_query}
        available_safe_asset_count = len(Task.objects.find(safe_ava_q).distinct("jobId"))
        return sec_risk_asset_count, available_risk_asset_count, available_safe_asset_count

    def asset_risk_data(self):
        """资产威胁等级计算"""
        q = {
            "uid": self.user.id,
            "endTime": self.time_query,
            "status": "completed",
            "taskType": {"$in": self.web_sec_type}
        }
        fields = {
            "taskType": 1,
            "result.warnings.level": 1,
            "result.warnings.title": 1,
            "result.warnings.affects": 1,
            "result.warnings.id": 1,
            "endTime": 1,
            "jobId": 1,
            "result.targetStatus": 1,
        }
        job_tasks_map = {}
        for t in Task._get_collection().find(q, fields):
                job_tasks_map.setdefault(t["jobId"], []).append(t)
        job_score_maps = []
        critical, high, medium, low, safe = 0, 0, 0, 0, 0
        top_risk_job = {}
        for job_id, tasks in job_tasks_map.items():
            score, level, risk_names = self.parse_job_level(tasks)
            job_score_maps.append({"job_id": job_id, "score": score, "risk_names": risk_names, "level": level})
            if level == "严重":
                critical += 1
            elif level == "高危":
                high += 1
            elif level == "中危":
                medium += 1
            elif level == "低危":
                low += 1
            elif level == "安全":
                safe += 1
        if job_score_maps:
            top_risk_job = list(sorted(job_score_maps, key=lambda x: x["score"], reverse=True))[0]
            job = Job.objects.filter(id=top_risk_job["job_id"]).first()
            top_risk_job.update({"name": job.note, "target_url": job.targetUrl})
        return critical, high, medium, low, safe, top_risk_job

    @staticmethod
    def parse_job_level(tasks):
        """资产安全"""
        all_warnings = []
        target_status_list = []
        end_times = []
        level = "未知"
        score = 0
        for task in tasks:
            if target_status := task.get("result", {}).get("targetStatus"):
                target_status_list.append(target_status.get("status"))
            else:
                target_status_list.append("good")
            all_warnings.extend(task.get("result", {}).get("warnings", []))
            end_times.append(task.get("endTime"))
        update_time = max(end_times) if end_times else ''
        all_warnings = list({w["id"]: w for w in all_warnings}.values())
        if all_warnings:
            score_dict = Score.score_info(event_map={}, warnings=all_warnings)
            score = score_dict.get('score', 0)
            level = score_dict.get('level', "未知")
        elif (not all_warnings and update_time) and (not target_status_list or "good" in target_status_list):
            level = '安全'
        elif (not all_warnings and update_time) and (target_status_list and "good" not in target_status_list):
            level = "未知"
        warning_titles = {i["title"]: i["level"] for i in all_warnings}
        warning_titles = [i[0] for i in sorted(warning_titles.items(), key=lambda x:x[1], reverse=True)]
        if warning_titles and len(warning_titles) > 3:
            risk_names = warning_titles[:3]
        else:
            risk_names = warning_titles
        return score, level, risk_names

    @staticmethod
    def parse_task_settings(task_settings, task_type):
        if task_type == "ssl":
            plugin_categories = task_settings.get("collect", {}).get("pluginCategories", [])
            data = [WARNING_CATEGORY_KEY_LABEL_MAP[task_type][p].replace("SSL", "") for p in plugin_categories]
        elif task_type == "content":
            plugins = task_settings.get("collect", {}).get("plugins", [])
            data = []
            if "keyword" in plugins:
                data.extend([keyword_category_mapping[int(t)] for t in set(task_settings.get("collect", {}).get("customKeywordTypes", []) + task_settings.get("collect", {}).get("systemKeywordTypes", []))])
            if "privacy_disclosure" in plugins:
                data.extend([privacy_category_mapping[int(t)] for t in task_settings.get("collect", {}).get("privacyDisclosureTypes", [])])
        elif task_type == "vul":
            data = []
            vul = task_settings.get("collect", {}).get("vul", ["1", ])
            vul_type = task_settings.get("vulType", "type")
            if vul_type == "plugIn":
                data.extend(vul)
            else:
                custom_ids = []
                system_ids = []
                for i in vul:
                    try:
                        system_ids.append(int(i))
                    except ValueError:
                        custom_ids.append(i)
                if custom_ids:
                    data.extend(KbVuln._get_db()["custom_profile_template"].distinct("plugin_list", {"custom_profile_template_id": {"$in": custom_ids}}))
                if system_ids:
                    data.extend(KbVuln._get_db()["system_profile_template"].distinct("plugin_list", {"value": {"$in": system_ids}}))
            if data:
                data = KbVuln._get_db()["kb_vuln"].distinct("name.zh_cn", {'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}, 'is_deleted': False, "key": {"$in": list(set(data))}})
        else:
            plugins = set()
            for p in task_settings.get("collect", {}).get("plugins", []):
                if "foreign_links" in p:
                    plugins.add("foreign_links")
                else:
                    plugins.add(p)
            data = [WARNING_CATEGORY_KEY_LABEL_MAP[task_type][p] for p in plugins]
        return data

    def web_sec_task_data(self):
        """web安全监测数据"""
        q = {
            "uid": self.user.id,
            "endTime": self.time_query,
            "status": "completed",
            "taskType": {"$in": self.web_sec_type}
        }
        fields = {
            "taskType": 1,
            "result.warnings.level": 1,
            "result.warnings.title": 1,
            "result.warnings.affects": 1,
            "result.warnings.id": 1,
            "taskSettings": 1,
            "result.targetStatus.status": 1,
            "taskSessionId": 1,
            "jobId": 1
        }
        task_session_ids = set()
        job_ids = set()
        tasks_map = {t: [] for t in self.web_sec_type}
        for t in Task._get_collection().find(q, fields):
            job_ids.add(t.get("jobId"))
            task_session_ids.add(t.get("taskSessionId"))
            tasks_map[t["taskType"]].append(t)
        source_data_map = {}
        completed_task_type_map = {t: False for t in self.web_sec_type}
        completed_task_count_map = {t: 0 for t in self.web_sec_type}
        good_sec_task_map = {t: False for t in self.web_sec_type}
        error_sec_task_count = 0
        settings_map = {t: [] for t in self.web_sec_type}
        settings_map.update({f"{t}_count": 0 for t in self.web_sec_type})
        for task_type, tasks in tasks_map.items():
            if not tasks:
                continue
            completed_task_type_map.update({task_type: True})
            for task in tasks:
                completed_task_count_map[task_type] += 1
                settings_map[task_type].extend(self.parse_task_settings(task["taskSettings"], task_type))
                for warning in task["result"]["warnings"]:
                    source_data_map.setdefault(task_type, {}).setdefault(
                        warning["level"], set()).add((warning["title"], warning["affects"], warning["id"]))
                if task["result"].get("targetStatus", {}).get("status", "good") == "good":
                    good_sec_task_map.update({task_type: True})
                else:
                    error_sec_task_count += 1
        risk_count = 0
        result_map = {}
        for task_type, risks in source_data_map.items():
            for level, risk_names in risks.items():
                result_map.setdefault(task_type, {}).setdefault(self.level_num_en[level], set()).update(risk_names)
                result_map.setdefault(task_type, {}).update({f"{self.level_num_en[level]}_count": 0})
                risk_count += len(risk_names)
        content_setting_map = {
            "keyword": [], "keyword_count": 0,
            "privacy": [], "privacy_count": 0,
        }
        for t, data in settings_map.items():
            if isinstance(data, int):
                continue
            unique_d = list(set(data))
            settings_map[t] = unique_d[:1] if unique_d else []
            settings_map[f"{t}_count"] = len(unique_d)
            if t == "content":
                if keyword := list(set(keyword_category_mapping.values()) & set(unique_d)):
                    keyword = list(sorted(keyword, reverse=True))
                    content_setting_map.update({"keyword": keyword[:1], "keyword_count": len(keyword)})
                if privacy := list(set(privacy_category_mapping.values()) & set(unique_d)):
                    content_setting_map.update({"privacy": privacy[:1], "privacy_count": len(privacy)})
        settings_map.update(content_setting_map)
        for t, risk_map in result_map.items():
            for level, risks in risk_map.items():
                if isinstance(risks, int):
                    continue
                risk_names = [i[0] for i in risks]
                risk_map[level] = list(set(risk_names))
                risk_map[f"{level}_count"] = len(risk_names)
        # 安全事件同一个类型有多个威胁等级， 展示高威胁等级
        if security_event_data := result_map.get("securityEvent", {}):
            critical = security_event_data.get("critical", [])
            high = security_event_data.get("high", [])
            medium = security_event_data.get("medium", [])
            low = security_event_data.get("low", [])
            if high:
                security_event_data["high"] = list(set(high) - set(critical))
            if medium:
                security_event_data["medium"] = list(set(medium) - (set(critical) | set(high)))
            if low:
                security_event_data["low"] = list(set(low) - (set(critical) | set(high) | set(medium)))
        web_sec_job_count = len(job_ids)
        scan_count = get_monitor_urls(task_session_id=list(task_session_ids), is_scan_count=True)
        return (result_map, risk_count, settings_map, completed_task_type_map,
                completed_task_count_map, good_sec_task_map, error_sec_task_count,
                web_sec_job_count, scan_count)

    def web_available_task_data(self):
        """可用性监测数据"""
        q = {
            "uid": self.user.id,
            "endTime": self.time_query,
            "status": "completed",
            "taskType": {"$in": self.available_type}
        }
        fields = {
            "taskType": 1,
            "result.addition.respAvg": 1,
            "result.addition.nodeTotalCount": 1,
            "result.addition.nodeAbnormalCount": 1,
            "jobId": 1,
            "endTime": 1
        }
        tasks_map = {}
        for t in Task._get_collection().find(q, fields):
            result = t.pop("result", {})
            t["respAvg"] = result.get("addition", {}).get("respAvg", 0)
            t["nodeTotalCount"] = result.get("addition", {}).get("nodeTotalCount", 0)
            t["nodeAbnormalCount"] = result.get("addition", {}).get("nodeAbnormalCount", 0)
            tasks_map.setdefault((t["jobId"], t["taskType"]), []).append(t)
        tmp_data = {}
        completed_task_type_map = {t: False for t in self.available_type}
        completed_task_count_map = {t: 0 for t in self.available_type}
        for _, tasks in tasks_map.items():
            if not tasks:
                continue
            job_id, task_type = _[0], _[1]
            completed_task_type_map.update({task_type: True})
            resp_avg_list = []
            ava_avg_list = []
            date_avg_map = {}
            for task in tasks:
                completed_task_count_map[task_type] += 1
                resp_avg_list.append(task["respAvg"])
                ava_avg_list.append(1 - round(task["nodeAbnormalCount"]/task["nodeTotalCount"], 2))
                date_str = self.utc_to_local(task["endTime"]).strftime("%Y-%m-%d")
                date_avg_map.setdefault(date_str, []).append(1 - round(task["nodeAbnormalCount"]/task["nodeTotalCount"], 2))
            resp_avg = 0
            if resp_avg_list:
                resp_avg = int(sum(resp_avg_list) / len(resp_avg_list))
            ava_avg = 0
            if ava_avg_list:
                ava_avg = round(sum(ava_avg_list) / len(ava_avg_list), 2)
            date_list = []
            for i in range((self.utc_to_local(self.end_time) - self.utc_to_local(self.start_time)).days + 1):
                date_list.append((self.utc_to_local(self.start_time) + datetime.timedelta(days=i)).strftime("%Y-%m-%d"))
            date_avg_list = [0, ] * len(date_list)
            for index, date_str in enumerate(date_list):
                if _ := date_avg_map.get(date_str):
                    date_avg_list[index] = round(sum(_) / len(_), 2)
                else:
                    # 空字符串 前端不会画点
                    date_avg_list[index] = ""
            tmp_data.setdefault(task_type, []).append({
                "job_id": job_id,
                "resp_avg": resp_avg,
                "ava_avg": ava_avg,
                "draw_data": {"date_list": date_list, "date_avg_list": date_avg_list}
            })
        res_data = []
        if tmp_data.get("http"):
            res_data = tmp_data["http"]
        elif tmp_data.get("ping"):
            res_data = tmp_data["ping"]
        res_data = list(sorted(res_data, key=lambda x: x["ava_avg"]))
        if len(res_data) > 3:
            res_data = res_data[:3]
        job_ids = [item.get("job_id") for item in res_data]
        job_data_map = {job.id: {"name": job.note, "target_url": job.targetUrl} for job in Job.objects.filter(id__in=job_ids)}
        for _ in res_data:
            job_id = _.pop("job_id")
            _.update(job_data_map.get(job_id, {}))
        return res_data, completed_task_type_map, completed_task_count_map

    def task_open_status(self):
        """监测项开启情况"""
        enable_task_map = {t: False for t in self.web_sec_type + self.available_type}
        if not self.is_free_user:
            monitor_asset_query = {
                "uid": self.user.id,
                "taskGroupType": {"$in": ['web_security', 'http', 'ping']},
                "$or": [
                    {"enableMonitor": True, "triggerType": "schedule"},
                    {"triggerType": "manual", "createTime": self.time_query}
                ]
            }
            for at in AssetTasks.objects.find(monitor_asset_query):
                if at.taskGroupType == "http":
                    enable_task_map["http"] = True
                elif at.taskGroupType == "ping":
                    enable_task_map["http"] = True
                else:
                    for t in self.web_sec_type:
                        if getattr(at, f"{t}Settings").enable:
                            enable_task_map[t] = True
        # 可能之前执行过任务， 但是现在关闭了开关
        completed_task_types = Task.objects.find({
            "uid": self.user.id,
            "taskType": {"$in": self.available_type + self.web_sec_type},
            "endTime": self.time_query,
            "status": "completed"
        }).distinct("taskType")
        # 场景3 当前开启监测， 但是统计周期没有任务执行完成， 不推送
        # 场景5 当前开启监测， 但是统计周期未开启（没办法得到统计周期是否开启了监测， 采用没有任务执行完成做依据）， 不推送
        is_push_report = False if any(enable_task_map.values()) and not completed_task_types else True
        for t in completed_task_types:
            enable_task_map.update({t: True})
        # 是否开启有监测
        enable_task_status = "null"
        if all(enable_task_map.values()):
            enable_task_status = "all"
        elif any(enable_task_map.values()):
            enable_task_status = "part"
        return is_push_report, enable_task_status, enable_task_map

    def get_data(self):
        is_push_report, enable_task_status, enable_task_map = self.task_open_status()
        if not is_push_report:
            return is_push_report, enable_task_map, "enable task but not run"
        receiver_data = self.get_receiver()
        result = {
            "uid": self.user.id,
            "username": self.user.username,
            "is_free_user": self.is_free_user,
            "start_time": self.utc_to_local(self.start_time).strftime("%Y/%m/%d"),
            "end_time": self.utc_to_local(self.end_time).strftime("%Y/%m/%d"),
            "enable_task_status": enable_task_status,
            "enable_task_map": enable_task_map,
            "receiver_data": receiver_data,
            "kb_data": self.kb_data,
            "push_time": arrow.utcnow().datetime
        }
        used_domain_count, free_count, expire_time = self.pkg_data()
        monitor_asset_count, not_monitor_asset_count, not_verified_asset_count = self.asset_data()
        sec_risk_asset_count, available_risk_asset_count, available_safe_asset_count = self.asset_task_data()
        result.update({
            "asset_status_data": {
                "used_domain_count": used_domain_count,
                "free_count": free_count,
                "expire_time": expire_time,
                "duration_time": self.duration_time,
                "monitor_asset_count": monitor_asset_count,
                "not_monitor_asset_count": not_monitor_asset_count,
                "not_verified_asset_count": not_verified_asset_count,
                "sec_risk_asset_count": sec_risk_asset_count,
                "available_risk_asset_count": available_risk_asset_count,
                "available_safe_asset_count": available_safe_asset_count
            }
        })
        critical, high, medium, low, safe, top_risk_job = self.asset_risk_data()
        result.update({
            "asset_risk_data": {
                "critical": critical,
                "high": high,
                "medium": medium,
                "low": low,
                "safe": safe,
                "top_risk_job": top_risk_job
            }
        })
        task_type_map, risk_count, settings_map, sec_completed_task_type, sec_completed_task_count, \
            good_sec_task_map, error_sec_task_count, web_sec_job_count, web_sec_scan_count = self.web_sec_task_data()
        result.update({
            "web_sec_task_data": {
                "task_type_map": task_type_map,
                "risk_count": risk_count,
                "good_sec_task_map": good_sec_task_map,
                "error_sec_task_count": error_sec_task_count,
                "web_sec_job_count": web_sec_job_count,
                "web_sec_scan_count": web_sec_scan_count,
            }
        })
        result["asset_status_data"].update({"web_sec_task_count": sum(sec_completed_task_count.values())})
        web_available_task_data, ava_completed_task_type, ava_completed_task_count = self.web_available_task_data()
        result.update({
            "web_available_task_data": web_available_task_data
        })
        result["asset_status_data"].update({"available_task_count": sum(ava_completed_task_count.values())})
        completed_task_type = {t: False for t in self.web_sec_type + self.available_type}
        completed_task_type.update(ava_completed_task_type)
        completed_task_type.update(sec_completed_task_type)
        completed_task_count = {t: 0 for t in self.web_sec_type + self.available_type}
        completed_task_count.update(ava_completed_task_count)
        completed_task_count.update(sec_completed_task_count)
        result.update({
            "overview_data": {
                "monitor_count": sum(completed_task_count.values()),
                "completed_task_count": completed_task_count,
                **settings_map,
            },
            "completed_task_type": completed_task_type
        })
        balance_enable_task_type = {t: False for t in self.web_sec_type + self.available_type}
        if self.balance:
            for t in (self.web_sec_type + self.available_type):
                if getattr(self.balance.feature, f"{t}Task"):
                    balance_enable_task_type.update({t: True})
        balance_enable_task_type.update({t: True for t, _ in completed_task_type.items() if _})
        enable_web_sec_balance = any([t for t in self.web_sec_type if balance_enable_task_type.get(t)])
        enable_web_ava_balance = any([t for t in self.available_type if balance_enable_task_type.get(t)])
        # 场景7 未开通web安全监测与可用性监测， 不推送
        if (not enable_web_sec_balance) and (not enable_web_ava_balance):
            is_push_report = False
        result["overview_data"].update({
            "enable_balance": {
                "web_sec": enable_web_sec_balance,
                "web_ava": enable_web_ava_balance
            }})
        current_app.logger.info(f"{self.user.username} report data all completed")
        return is_push_report, result, "success" if is_push_report else "fail"
